Skip to content

feat(libtorch): report device identity per index - #35

Open
NicolasRouquette wants to merge 5 commits into
lean-dojo:mainfrom
NicolasRouquette:libtorch-device-identity
Open

NicolasRouquette wants to merge 5 commits into
lean-dojo:mainfrom
NicolasRouquette:libtorch-device-identity

Conversation

@NicolasRouquette

@NicolasRouquette NicolasRouquette commented Oct 3, 2026 •

Copy link
Copy Markdown
Contributor

Stacked on #36

This branch carries #36's three commits first; this PR's own are the last two. #36 gives the
bridge its host device, and this PR's queries must say what they do there, so the two are checked
together. Merging #36 first leaves this PR with its own two commits.

What

Runtime.Autograd.LibTorch.deviceInfo (index : UInt32) : IO DeviceInfo and
currentDeviceInfo, the LibTorch follow-up to #30: name, compute capability (major * 10 + minor, so 86 reads as sm_86), SM count, core and memory clocks, memory bus width, total
memory, and the driver and runtime versions, plus DeviceInfo.peakBandwidthGBs and a one-line
DeviceInfo.format for a benchmark header. Buffer.memory answers how much device memory there
is; this answers which device.

The review finding from #30

#30 cached one cudaDeviceProp under a process-wide pthread_once, so after switching devices
it kept reporting device 0 and could combine fields from two cards. Here every query takes an
explicit index and reads at::cuda::getDeviceProperties(index), the SDK's own per-device cache,
so the answer never depends on which device was selected earlier; currentDeviceInfo is
deviceInfo (← getDevice). The test switches through every visible device in a fresh process
and checks that currentDeviceInfo reports each index in turn.

Clocks and bus width come from cudaDeviceGetAttribute, since CUDA 13 removed clockRate and
memoryClockRate from cudaDeviceProp and the attribute enums exist in 12.x and 13.x. An index
at or above device_count is a TORCH_CHECK failure that surfaces as an IO error, like every
other runtime control; the build without LibTorch fails the same way through unavailable.c
rather than returning a record of zeroes. On a CPU-only SDK, where cudaDeviceProp does not
exist, the exports compile under TORCHLEAN_LIBTORCH_CUDA and fail through IO the same way. A
CUDA-enabled SDK answers even while the bridge runs on the host: the queries describe the visible
devices, not the selected one.

What is not carried over

#30 also reported the architectures the binary held native code for (__CUDA_ARCH_LIST__) and
warned when the device was being served through PTX JIT. Neither carries over, for two separate
reasons the module docstring keeps apart. The SDK's packaged architectures are a property of the
SDK build, and ATen's C++ API does not expose the list (only the Python package reads it back).
The kernels TorchLean compiles itself, the generated operations of NN.Kernel, are compiled at
run time by NVRTC for the compute capability of the device they run on, so there is no
compiled-in architecture list for them to report.

Tests

NN/Tests/Runtime/Cuda/DeviceInfo.lean, called from NN/Tests/Suite.lean under every runtime
status:

  • with a device: every visible index reports itself, a nonempty name, a nonzero capability, SM
    count, and memory; currentDeviceInfo agrees with getDevice; the index at the visible count
    is rejected; and a child process (TORCHLEAN_LIBTORCH_DEVICE_PROBE=switch, the pattern of the
    memory probes) walks setDevice through every device before any buffer exists, because
    setDevice refuses while wrappers are live;
  • with the bridge on the host and a CUDA-enabled SDK: the readback alone, since the host rejects
    setDevice; the switching probe requires a CUDA device (Buffer.requireCudaDevice, from feat(libtorch): run the bridge on the host when the SDK is CPU-only #36);
  • on a CPU-only SDK, and without LibTorch: deviceInfo 0 fails.

Checks run

RTX A4500 (sm_86, driver 580.126), pip torch 2.11.0+cu128 with a CUDA 12.8 toolkit for SDK
discovery and NVRTC, and pip torch 2.11.0+cpu with no toolkit; Lean v4.35.0, on cfbcb9a5:

lake -R -Kcuda=true build nn_tests_suite (either SDK)          Build completed successfully (8535 jobs).
TORCHLEAN_REQUIRE_CUDA=1 nn_tests_suite                        == TorchLean: all curated tests passed ==
  same binary, TORCHLEAN_LIBTORCH_DEVICE=cpu                   == TorchLean: all curated tests passed ==
  CPU-only SDK build, TORCHLEAN_REQUIRE_NATIVE=1               == TorchLean: all curated tests passed ==
lake -R -Kcuda=false build NN NNTests nn_tests_suite NNCI NNExamples   Build completed successfully (10204 jobs).
nn_tests_suite (default build)                                 device identity: queries are rejected without LibTorch
                                                               == TorchLean: all curated tests passed ==
python3 scripts/checks/repo_lint.py --fail-on-warn             OK: no issues found.

Device-identity section of the GPU run:

== LibTorch device identity ==
  NVIDIA RTX A4500 (device 0, sm_86, 56 SMs, 19.548828 GiB, ~640.080000 GB/s peak, driver 13.0, runtime 12.8)
  device switching: 1 device(s) report under their own index

On the host with the same binary the record is the same and the switching probe reads
skipped: the device switching probe needs the bridge on a CUDA device; on the CPU-only SDK,
device identity: queries are rejected on a CPU-only SDK.

Not run: the sanitizer harness and the elementwise C++ harness. This host has one GPU, so the
switching probe exercised one index; the multi-device case is covered by construction (per-index
reads) rather than by an observation.

NicolasRouquette added a commit to NicolasRouquette/TorchLean that referenced this pull request Oct 3, 2026
Brings in the upstream PR branch (lean-dojo#35) at 65fd9d9, one
commit on upstream main b062b9a: Runtime.Autograd.LibTorch.deviceInfo and
currentDeviceInfo read at::cuda::getDeviceProperties(index) and
cudaDeviceGetAttribute per index, with DeviceInfo.format for benchmark
headers.

Successor of the closed lean-dojo#30. The compiled-architecture fields are gone with
the custom kernels.
@Robertboy18

Copy link
Copy Markdown
Member

Hey Nicolas, thanks! Per-device identity is useful. Could you check this together with #36? The CUDA property queries need CPU-only SDK handling, and the tests should not assume nativeAvailable means a CUDA device exists. Also, the statement in DeviceInfo.lean that TorchLean compiles no device code of its own will need updating when our pending custom-kernel/NVRTC work lands. Please keep the SDK-packaged architecture list separate from runtime-compiled kernels in that explanation. Keeping this open for those integration checks.

The LibTorch bridge pinned `c10::kCUDA`, so a build without a CUDA device had no backend at all:
`unavailable.c` fails every buffer operation, and the device buffer suite could only run on a
GPU host. The bridge now selects its device once, at its first call, and runs on the host through
ATen's CPU kernels when the SDK has no CUDA support or when `TORCHLEAN_LIBTORCH_DEVICE=cpu` asks
for it. Nothing is reimplemented: every ATen operation already goes through `torchlean::options()`.

Selection: `TORCHLEAN_LIBTORCH_DEVICE=cpu` selects the host and `=cuda` insists on a CUDA device.
Unset, a visible CUDA device is selected when the SDK has CUDA support; otherwise the host is
selected only when the SDK has no CUDA support at all, so a CUDA-enabled SDK without a visible
device stays `RuntimeStatus.nativeUnavailable` rather than silently running on the host.

CMake no longer requires the `torch_cuda` target. A CPU-only SDK compiles the backend with
`TORCHLEAN_LIBTORCH_CUDA=0`, without CUDA headers, the CUDA toolkit, NVRTC or the driver API; the
CPU-only pip wheel is a complete SDK for it. `sdk.txt` records the choice.

Runtime-compiled kernels (the generated operations of `NN.Kernel`) need NVRTC and the driver API,
so they compile only with a CUDA-enabled SDK, and `run_kernel` refuses at its entry unless a CUDA
device is selected: on the host a generated operation fails through `IO` with an error that names
the requirement, before any CUDA call. The CUDA suite runs its custom-computation checks on a CUDA
device and checks that refusal on the host.

Lean: `Runtime.Autograd.LibTorch.DeviceKind` and `deviceKind` report the selection.
`RuntimeStatus` keeps its three states, with `.nativeAvailable` meaning that a device is
selected, CUDA or host, so `requireNativeRuntime` and every existing match are unchanged. On the
host the native allocator counters and driver memory read zero, `synchronize` and `emptyCache`
return at once, `setMemoryFraction` and `setDevice` are rejected, and the suite skips its native
memory probes. The suite prints the device and checks that the selection agrees with the visible
device count and with the environment variable.
…nd CUDA

Programs made of the IEEE basic operations, gathers and lookups give bit-identical results on
the two devices; transcendental functions come from different libraries on the host and on
CUDA and differ at the ulp level, and reductions differ in order.
`nativeAvailable` covers the host now that the bridge can run there, so `requireNativeRuntime`,
which the suite called for `TORCHLEAN_REQUIRE_CUDA=1`, let a run on a CPU-only SDK or under
`TORCHLEAN_LIBTORCH_DEVICE=cpu` pass a check meant to say that a GPU ran it.
`Buffer.requireCudaDevice` asks for the device kind as well, and the suite uses it for that
variable. `TORCHLEAN_REQUIRE_NATIVE=1` keeps the bridge-only meaning for CPU-only SDK sessions,
so that such a job fails on the unlinked default backend instead of passing with the device suite
skipped.

In the bridge, `require_cuda("<feature>")` is the entry check for an operation with no host
implementation (device selection, the allocator memory fraction and runtime-compiled kernels use
it), so it fails with a message naming the feature rather than with a CUDA call on a CPU tensor. The
READMEs, the installation page and the execution-modes chapter say what each variable rejects, and
quote the new messages.

Configured binary arithmetic computes only in runtime-compiled kernels, which the tape reaches
through pure entry points, where an error ends the process. The encoded upload therefore calls
`require_cuda` too, so on the host such a computation fails through `IO` when its tensor is
uploaded. The host branch of the CUDA suite checks both refusals and runs the custom-autograd checks
that need no generated kernel: the native-precision recordings and the resident-buffer checks.
`LibTorch.deviceInfo index` reads one device's name, compute capability, SM count, clocks, bus
width, total memory, and the driver and runtime versions through the SDK's per-device cached
properties; `currentDeviceInfo` reads the device `getDevice` selects. Querying by explicit index
keeps the answer independent of the selected device, so switching devices cannot report a stale
record. Without LibTorch, the queries fail through `IO` like the other controls.

`NN/Tests/Runtime/Cuda/DeviceInfo.lean` checks every visible device under its own index, the
current-device agreement, index rejection, and, in a fresh process selected by
`TORCHLEAN_LIBTORCH_DEVICE_PROBE=switch`, that switching devices changes the reading.
…ty queries

The identity exports read `cudaDeviceProp` and `cudaDeviceGetAttribute`, which a CPU-only SDK
does not provide, and the test read `nativeAvailable` as "a CUDA device is present", which the
host bridge no longer implies. The exports now compile under `TORCHLEAN_LIBTORCH_CUDA`; a CPU-only
SDK gets exports that fail through `IO`, as an index outside the visible count does on a
CUDA-enabled SDK. A CUDA-enabled SDK still answers while the bridge runs on the host, since the
queries describe the visible devices rather than the selected one.

`NN/Tests/Runtime/Cuda/DeviceInfo.lean` follows the device the bridge selected: on a CUDA device
it runs the readback and the switching probe, which now requires a CUDA device; on the host with
a CUDA-enabled SDK it runs the readback alone, because the host rejects `setDevice`; on a
CPU-only SDK it checks that every query fails.

The module docstring keeps the SDK's packaged architectures, a property of the SDK build, apart
from the kernels TorchLean compiles at run time: NVRTC compiles the generated operations of
`NN.Kernel` for the compute capability of the device they run on.
@NicolasRouquette
NicolasRouquette force-pushed the libtorch-device-identity branch from 65fd9d9 to cfbcb9a Compare October 19, 2026 01:06
@NicolasRouquette

Copy link
Copy Markdown
Contributor Author

Rebased onto current main (72754452, Lean 4.35) and, since you asked for it to be checked
together with #36, stacked on #36: the first three commits are #36's, and this PR's own are the
last two, d7f5e85c (the change) and cfbcb9a5 (the review commit). Merging #36 first leaves
this PR with those two.

  • CPU-only SDK. The identity exports read cudaDeviceProp and cudaDeviceGetAttribute,
    which a CPU-only SDK does not provide, so they now compile under TORCHLEAN_LIBTORCH_CUDA; a
    CPU-only SDK gets exports that fail through IO (LibTorch: this SDK has no CUDA support, so there is no CUDA device identity to read), as an index outside the visible count fails on a
    CUDA-enabled SDK. A CUDA-enabled SDK keeps answering while the bridge runs on the host, since the
    queries describe the visible devices rather than the selected one.
  • nativeAvailable is no longer read as a CUDA device. NN/Tests/Runtime/Cuda/DeviceInfo.lean
    follows the device the bridge selected: on a CUDA device, the readback and the switching probe,
    which now calls feat(libtorch): run the bridge on the host when the SDK is CPU-only #36's requireCudaDevice; on the host with a CUDA-enabled SDK, the readback
    alone, because the host rejects setDevice; on a CPU-only SDK, a check that every query fails.
  • DeviceInfo.lean and the NVRTC work. Your custom-operation path is on main now, so the
    docstring no longer says TorchLean compiles no device code. It keeps the two questions apart, as
    you asked: the SDK's packaged GPU architectures, the sm_XY targets its own kernels were built
    for, are a property of the SDK build that ATen's C++ API does not expose; the kernels TorchLean
    compiles at run time are compiled by NVRTC for the compute capability of the device they run on
    (compile_kernel passes --gpu-architecture=compute_XY from the current device's properties),
    whatever the SDK was packaged for, so their coverage is a question for NVRTC, not for this
    readback.
  • What else the rebase changed. main renamed Buffer.allocatorStats to Buffer.memory, and
    the docstring follows; the suite calls Tests.Cuda.DeviceInfo.run right after
    Tests.Floats.run, since main dropped the rationals suite it used to follow.

Verification. RTX A4500, pip torch 2.11.0+cu128 with a CUDA 12.8 toolkit, and pip torch
2.11.0+cpu with no toolkit; Lean v4.35.0, on cfbcb9a5. The device-identity section of each run:

GPU (TORCHLEAN_REQUIRE_CUDA=1):
  == LibTorch device identity ==
    NVIDIA RTX A4500 (device 0, sm_86, 56 SMs, 19.548828 GiB, ~640.080000 GB/s peak, driver 13.0, runtime 12.8)
    device switching: 1 device(s) report under their own index
    LibTorch device: cuda
same binary on the host (TORCHLEAN_LIBTORCH_DEVICE=cpu TORCHLEAN_REQUIRE_NATIVE=1):
  == LibTorch device identity ==
    NVIDIA RTX A4500 (device 0, sm_86, 56 SMs, 19.548828 GiB, ~640.080000 GB/s peak, driver 13.0, runtime 12.8)
    skipped: the device switching probe needs the bridge on a CUDA device
    LibTorch device: host
CPU-only wheel, no toolkit (TORCHLEAN_REQUIRE_NATIVE=1):
  == LibTorch device identity ==
    device identity: queries are rejected on a CPU-only SDK
    LibTorch device: host
default build:
  == LibTorch device identity ==
    device identity: queries are rejected without LibTorch

All four runs end in == TorchLean: all curated tests passed ==, and the three strict runs of #36
(TORCHLEAN_REQUIRE_CUDA=1 on the host and on the CPU-only SDK, TORCHLEAN_REQUIRE_NATIVE=1 on
the default build) fail as they should. Builds: -K cuda=true 8535 jobs against either SDK, the
default build 10204. python3 scripts/checks/repo_lint.py --fail-on-warn: OK. The PR description
is updated to match.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants